import torch
import torchvision.models as models
import torch.quantization

# 加载一个预训练的ResNet模型
model = models.resnet18(pretrained=True)
model.eval()

# 准备一个输入张量用于量化配置
example_input = torch.rand(1, 3, 224, 224)

# 配置模型以进行动态量化
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
torch.quantization.prepare_qat(model, inplace=True)

# 使用量化感知训练对模型进行微调（这里仅作为示例，实际应用中需要完整的训练循环）
# model.train()
# optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# criterion = torch.nn.CrossEntropyLoss()
# for epoch in range(num_epochs):
#     # 训练循环...

# 转换模型为量化模型
torch.quantization.convert(model, inplace=True)

# 验证量化后的模型
model.eval()
with torch.no_grad():
    output = model(example_input)
    print(output)
